from detectron2.config import CfgNode as CN


def add_p3aformer_config(cfg, args):
    """
    Add config for MASK_FORMER.
    """
    # MOT dataset configs
    cfg.INPUT.DATASET_MAPPER_NAME = "mot_mixed"  # specific name
    cfg.INPUT.MOT_DATASET_NAME = "MOT17"  # which MOT dataset
    cfg.INPUT.DATA_DIR = "./data/mix_det"
    cfg.INPUT.VAL_DATA_DIR = "./data/mot"
    cfg.INPUT.SPLIT = "val_half"
    cfg.INPUT.COLOR_AUG_SSD = True
    cfg.SEED = 2022

    # p3aformer settings
    cfg.MODEL.DENSETRACK = CN()
    cfg.MODEL.DENSETRACK.NUM_FEATURE_LEVELS = 4
    cfg.MODEL.DENSETRACK.ENC_LAYERS = 6
    cfg.MODEL.DENSETRACK.DEC_LAYERS = 6
    cfg.MODEL.DENSETRACK.DIM_FEEDFORWARD = 1024
    cfg.MODEL.DENSETRACK.HIDDEN_DIM = 256
    cfg.MODEL.DENSETRACK.NHEADS = 8
    cfg.MODEL.DENSETRACK.POSITION_EMBEDDING = "sine"
    cfg.MODEL.DENSETRACK.MASKS = False
    cfg.MODEL.DENSETRACK.BACKBONE = "resnet50"
    cfg.MODEL.DENSETRACK.DILATION = False
    cfg.MODEL.DENSETRACK.DROPOUT = 0.1
    cfg.MODEL.DENSETRACK.DEC_N_POINTS = 4
    cfg.MODEL.DENSETRACK.ENC_N_POINTS = 4
    cfg.MODEL.DENSETRACK.TRACKING = True
    cfg.MODEL.DENSETRACK.DOWN_RATIO = 4  # image size
    cfg.MODEL.DENSETRACK.SAME_AUG_PRE = True
    cfg.MODEL.DENSETRACK.PRE_HM = True
    cfg.MODEL.DENSETRACK.HEADS = ["hm", "reg", "wh", "center_offset", "tracking"]
    cfg.MODEL.DENSETRACK.HM_WEIGHT = 1.0
    cfg.MODEL.DENSETRACK.OFF_WEIGHT = 1.0
    cfg.MODEL.DENSETRACK.WH_WEIGHT = 0.1
    cfg.MODEL.DENSETRACK.BOXES_WEIGHT = 0.5
    cfg.MODEL.DENSETRACK.GIOU_WEIGHT = 0.4
    cfg.MODEL.DENSETRACK.CT_OFFSET_WEIGHT = 0.1
    cfg.MODEL.DENSETRACK.TRACKING_WEIGHT = 1.0
    cfg.MODEL.DENSETRACK.NORM_FACTOR = 1.0
    cfg.MODEL.DENSETRACK.DEFAULT_RESOLUTION = [640, 1088]

    # Tracking settings
    cfg.TRACK = CN()
    cfg.TRACK.VIS = False  # open visualization
    cfg.TRACK.DENSETRACK = CN()
    cfg.TRACK.DENSETRACK.PRIVATE = True
    cfg.TRACK.DENSETRACK.TRACK_THRE = 0.65
    cfg.TRACK.DENSETRACK.LOW_THRE = 0.2
    cfg.TRACK.DENSETRACK.FIRST_ASSIGN_THRE = 0.9
    cfg.TRACK.DENSETRACK.SECOND_ASSIGN_THRE = 0.5
    cfg.TRACK.DENSETRACK.PRE_THRE = 0.5
    cfg.TRACK.DENSETRACK.OUT_THRE = 0.3
    cfg.TRACK.DENSETRACK.K_ELEMENT = 300

    # Output Settings
    cfg.OUTPUT_DIR = "output/debug"
    # # Color augmentation
    # cfg.INPUT.COLOR_AUG_SSD = False
    # # We retry random cropping until no single category in semantic segmentation GT occupies more
    # # than `SINGLE_CATEGORY_MAX_AREA` part of the crop.
    # cfg.INPUT.CROP.SINGLE_CATEGORY_MAX_AREA = 1.0
    # # Pad image and segmentation GT in dataset mapper.
    # cfg.INPUT.SIZE_DIVISIBILITY = -1

    # solver config
    # weight decay on embedding
    cfg.SOLVER.WEIGHT_DECAY_EMBED = 0.0
    cfg.SOLVER.OPTIMIZER = "ADAMW"
    cfg.SOLVER.BACKBONE_MULTIPLIER = 0.1
    cfg.SOLVER.LR_BACKBONE = 2e-5
    cfg.SOLVER.AUX_LOSS = False
    if args.eval_only:
        cfg.SOLVER.EVAL = True
    else:
        cfg.SOLVER.EVAL = False

    # # mask_former model config
    # cfg.MODEL.MASK_FORMER = CN()

    # # loss
    # cfg.MODEL.MASK_FORMER.DEEP_SUPERVISION = True
    # cfg.MODEL.MASK_FORMER.NO_OBJECT_WEIGHT = 0.1
    # cfg.MODEL.MASK_FORMER.CLASS_WEIGHT = 1.0
    # cfg.MODEL.MASK_FORMER.DICE_WEIGHT = 1.0
    # cfg.MODEL.MASK_FORMER.MASK_WEIGHT = 20.0

    # # transformer config
    # cfg.MODEL.MASK_FORMER.NHEADS = 8
    # cfg.MODEL.MASK_FORMER.DROPOUT = 0.1
    # cfg.MODEL.MASK_FORMER.DIM_FEEDFORWARD = 2048
    # cfg.MODEL.MASK_FORMER.ENC_LAYERS = 0
    # cfg.MODEL.MASK_FORMER.DEC_LAYERS = 6
    # cfg.MODEL.MASK_FORMER.PRE_NORM = False

    # cfg.MODEL.MASK_FORMER.HIDDEN_DIM = 256
    # cfg.MODEL.MASK_FORMER.NUM_OBJECT_QUERIES = 100

    # cfg.MODEL.MASK_FORMER.TRANSFORMER_IN_FEATURE = "res5"
    # cfg.MODEL.MASK_FORMER.ENFORCE_INPUT_PROJ = False

    # # mask_former inference config
    # cfg.MODEL.MASK_FORMER.TEST = CN()
    # cfg.MODEL.MASK_FORMER.TEST.SEMANTIC_ON = True
    # cfg.MODEL.MASK_FORMER.TEST.INSTANCE_ON = False
    # cfg.MODEL.MASK_FORMER.TEST.PANOPTIC_ON = False
    # cfg.MODEL.MASK_FORMER.TEST.OBJECT_MASK_THRESHOLD = 0.0
    # cfg.MODEL.MASK_FORMER.TEST.OVERLAP_THRESHOLD = 0.0
    # cfg.MODEL.MASK_FORMER.TEST.SEM_SEG_POSTPROCESSING_BEFORE_INFERENCE = False

    # # Sometimes `backbone.size_divisibility` is set to 0 for some backbone (e.g. ResNet)
    # # you can use this config to override
    # cfg.MODEL.MASK_FORMER.SIZE_DIVISIBILITY = 32

    # # pixel decoder config
    # cfg.MODEL.SEM_SEG_HEAD.MASK_DIM = 256
    # # adding transformer in pixel decoder
    # cfg.MODEL.SEM_SEG_HEAD.TRANSFORMER_ENC_LAYERS = 0
    # # pixel decoder
    # cfg.MODEL.SEM_SEG_HEAD.PIXEL_DECODER_NAME = "BasePixelDecoder"

    # # swin transformer backbone
    # cfg.MODEL.SWIN = CN()
    # cfg.MODEL.SWIN.PRETRAIN_IMG_SIZE = 224
    # cfg.MODEL.SWIN.PATCH_SIZE = 4
    # cfg.MODEL.SWIN.EMBED_DIM = 96
    # cfg.MODEL.SWIN.DEPTHS = [2, 2, 6, 2]
    # cfg.MODEL.SWIN.NUM_HEADS = [3, 6, 12, 24]
    # cfg.MODEL.SWIN.WINDOW_SIZE = 7
    # cfg.MODEL.SWIN.MLP_RATIO = 4.0
    # cfg.MODEL.SWIN.QKV_BIAS = True
    # cfg.MODEL.SWIN.QK_SCALE = None
    # cfg.MODEL.SWIN.DROP_RATE = 0.0
    # cfg.MODEL.SWIN.ATTN_DROP_RATE = 0.0
    # cfg.MODEL.SWIN.DROP_PATH_RATE = 0.3
    # cfg.MODEL.SWIN.APE = False
    # cfg.MODEL.SWIN.PATCH_NORM = True
    # cfg.MODEL.SWIN.OUT_FEATURES = ["res2", "res3", "res4", "res5"]
    # cfg.MODEL.SWIN.USE_CHECKPOINT = False

    # # NOTE: maskformer2 extra configs
    # # transformer module
    # cfg.MODEL.MASK_FORMER.TRANSFORMER_DECODER_NAME = "MultiScaleMaskedTransformerDecoder"

    # # LSJ aug
    # cfg.INPUT.IMAGE_SIZE = 1024
    # cfg.INPUT.MIN_SCALE = 0.1
    # cfg.INPUT.MAX_SCALE = 2.0

    # # MSDeformAttn encoder configs
    # cfg.MODEL.SEM_SEG_HEAD.DEFORMABLE_TRANSFORMER_ENCODER_IN_FEATURES = ["res3", "res4", "res5"]
    # cfg.MODEL.SEM_SEG_HEAD.DEFORMABLE_TRANSFORMER_ENCODER_N_POINTS = 4
    # cfg.MODEL.SEM_SEG_HEAD.DEFORMABLE_TRANSFORMER_ENCODER_N_HEADS = 8

    # # point loss configs
    # # Number of points sampled during training for a mask point head.
    # cfg.MODEL.MASK_FORMER.TRAIN_NUM_POINTS = 112 * 112
    # # Oversampling parameter for PointRend point sampling during training. Parameter `k` in the
    # # original paper.
    # cfg.MODEL.MASK_FORMER.OVERSAMPLE_RATIO = 3.0
    # # Importance sampling parameter for PointRend point sampling during training. Parametr `beta` in
    # # the original paper.
    # cfg.MODEL.MASK_FORMER.IMPORTANCE_SAMPLE_RATIO = 0.75
